# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import asyncio
import logging
import os
import signal
from concurrent.futures import ThreadPoolExecutor

import grpc
from tornado.ioloop import (IOLoop, PeriodicCallback)

from coin.base.param_util import to_list
from coin.util.queue.publisher import AsyncKafkaPublisher
from coin.util.queue.reader import AsyncKafkaReader
from coin.util.queue.config import KafkaConfig


class EngineBase(object):
  # User implement -----------------------------------------
  def __init__(self):
    """Child classes must do belows.
          (1) super().__init__()
          (2) Set topics using self.set_engine_config(topics=[...]).
        """
    self._engine_config = {}
    self._ioloop = None
    self._publisher = None
    self._grpc_server = None
    self._shutdown_func = None

  def prepare(self):
    """Invoked once all dependencies are ready, before the loop is started.
        """

  def shutdown(self):
    """Invoked before shutting down.
        """

  def on_message(self, topic, message):
    raise NotImplementError()

  def on_ping(self):
    pass

  def on_shutdown_requested(self):
    pass

  def on_publish_error(self, topic, message):
    pass

  def on_exception(self, exception):
    pass

  # Low-level user implement -------------------------------
  def on_kafka_records(self, records):
    for record in records:
      try:
        self.on_message(record.topic, record.value)
      except Exception as e:
        logging.exception('Exception raised from engine')
        self.on_exception(e)

  # API ----------------------------------------------------
  def publish(self, topic, *msgs):
    if self._publisher is not None:
      self._publisher.publish(topic, *msgs)
    else:
      logging.error('publish too early: %s', topic)

  def request_shutdown(self):
    self._shutdown_func()

  def set_engine_config(self,
                        *,
                        topics=None,
                        timestamp_from=None,
                        ping_period_ms=None,
                        grpc_enable=None,
                        grpc_num_workers=None,
                        grpc_address=None):
    if topics is not None:
      self._engine_config['topics'] = to_list(topics)

    if timestamp_from is not None:
      self._engine_config['timestamp_from'] = timestamp_from

    if ping_period_ms is not None:
      self._engine_config['ping_period_ms'] = ping_period_ms

    if grpc_enable is not None:
      self._engine_config['grpc_enable'] = grpc_enable

    if grpc_num_workers is not None:
      self._engine_config['grpc_num_workers'] = grpc_num_workers

    if grpc_address is not None:
      self._engine_config['grpc_address'] = grpc_address

  @property
  def ioloop(self):
    return self._ioloop

  @property
  def grpc_server(self):
    return self._grpc_server

  # Internal use -------------------------------------------
  def inject_deps(self, *, ioloop, publisher, grpc_server, shutdown_func):
    self._ioloop = ioloop
    self._publisher = publisher
    self._grpc_server = grpc_server
    self._shutdown_func = shutdown_func

  @property
  def engine_config(self):
    return self._engine_config


class _EngineIgnitor(object):
  def __init__(self, engine, kafka_config=None, catch_signal=True):
    self._kafka_config = kafka_config
    assert kafka_config is not None
    self._engine = engine
    self._ioloop = None
    self._reader = None
    self._publisher = None
    self._grpc_server = None

    self._catch_signal = catch_signal
    self._sig_counter = 0

    self._shutdown_requested = False

  def run_sync(self):
    if self._catch_signal:
      signal.signal(signal.SIGTERM, self._signal_handler)
      signal.signal(signal.SIGINT, self._signal_handler)

    self._ioloop = IOLoop.current()
    self._ioloop.run_sync(self._loop)

  def _request_shutdown(self):
    try:
      if not self._shutdown_requested:
        logging.warning('Shutdown requested. Will shutdown in 3 seconds.')
    finally:
      self._shutdown_requested = True

  def _signal_handler(self, signum, frame):
    self._sig_counter += 1
    if self._sig_counter == 4:
      logging.error('Exiting immediately.')
      os._exit(1)
    self._ioloop.add_callback_from_signal(self._request_shutdown)

  async def _loop(self):
    # Init
    await self._init_deps()
    await asyncio.coroutine(self._engine.prepare)()

    # Start
    if self._grpc_server:
      self._grpc_server.start()

    ping_callback = PeriodicCallback(self._engine.on_ping,
                                     self._engine.engine_config.get('ping_period_ms', 1000))
    ping_callback.start()

    # Loop
    while not self._shutdown_requested:
      try:
        records = await self._reader.read(timeout_ms=100)
      except TimeoutError:
        continue

      try:
        self._engine.on_kafka_records(records)
      except Exception as e:
        logging.exception('Exception raised from engine')
        self._engine.on_exception(e)

    # Shutdown
    await asyncio.coroutine(self._engine.on_shutdown_requested)()
    await asyncio.sleep(3.)

    ping_callback.stop()
    if self._grpc_server:
      self._grpc_server.stop(grace=3)

    await self._shutdown()

  async def _init_deps(self):
    self._publisher = AsyncKafkaPublisher(kafka_config=self._kafka_config,
                                          error_callback=self._engine.on_publish_error)
    self._publisher.start(ioloop=self._ioloop)

    self._reader = AsyncKafkaReader(self._engine.engine_config['topics'],
                                    kafka_config=self._kafka_config,
                                    timestamp_from=self._engine.engine_config.get(
                                        'timestamp_from', None))
    await self._reader.open()

    if self._engine.engine_config.get('grpc_enable', False):
      grpc_num_workers = self._engine.engine_config.get('grpc_num_workers', 1)
      grpc_address = self._engine.engine_config.get('grpc_address', '0.0.0.0:50051')
      self._grpc_server = grpc.server(ThreadPoolExecutor(max_workers=grpc_num_workers))
      self._grpc_server.add_insecure_port(grpc_address)

    self._engine.inject_deps(ioloop=self._ioloop,
                             publisher=self._publisher,
                             grpc_server=self._grpc_server,
                             shutdown_func=self._request_shutdown)

  async def _shutdown(self):
    await self._reader.close()
    await asyncio.coroutine(self._engine.shutdown)()
    self._publisher.stop()


def run_engine(engine, *, kafka_servers=None, kafka_config=None, catch_signal=True):
  if kafka_servers is not None:
    assert kafka_config is None
    kafka_config = KafkaConfig.from_kafka_servers(kafka_servers)

  runner = _EngineIgnitor(engine, kafka_config=kafka_config, catch_signal=catch_signal)
  runner.run_sync()
